import math

import numpy as np
import numpy.ma as ma

import matplotlib
rcParams = matplotlib.rcParams
from matplotlib.artist import kwdocd
from matplotlib.axes import Axes
from matplotlib import cbook
from matplotlib.patches import Circle
from matplotlib.path import Path
from matplotlib.ticker import Formatter, Locator, NullLocator, FixedLocator, NullFormatter
from matplotlib.transforms import Affine2D, Affine2DBase, Bbox, \
    BboxTransformTo, IdentityTransform, Transform, TransformWrapper


from coords import coord_conv


class AuxTransformAxes(Axes):

    def __init__(self, parent_axes, aux_transform):

        self.transAux = aux_transform

        self._parent_axes = parent_axes
        self.bbox = parent_axes.bbox

        Axes.__init__(self,
                      parent_axes.figure,
                      parent_axes._position, frameon=False)

        #self.cla()

    def _set_lim_and_transforms(self):

        Axes._set_lim_and_transforms(self)
        
        self.transAxes = self._parent_axes.transAxes
        self.transScale = self._parent_axes.transScale
        self.transLimits = self._parent_axes.transLimits

        self.transData = \
            self.transAux + \
            self._parent_axes.transData

#     def get_position(self, original=False):
#         return self._parent_axes.get_position(original)

#     def get_aspect(self):
#         return self._parent_axes.get_aspect()

#     def get_data_ratio(self):
#         return self._parent_axes.get_data_ratio()

#     def get_anchor(self):
#         return self._parent_axes.get_anchor()

    def apply_aspect(self, position=None):
        pos = self._parent_axes.get_position(original=False)
        self.set_position(pos, "active")


    def cla(self):
        super(AuxTransformAxes, self).cla()

        self.set_xticks([])
        self.set_yticks([])


    def update_datalim(self, *kl, **kw):
        pass

    def set_xlim(self, *kl, **kw):
        pass

    def set_ylim(self, *kl, **kw):
        pass



class WcsSky2PixelTransform(Transform):
    """
    The base Aitoff transform.
    """
    input_dims = 2
    output_dims = 2
    is_separable = False

    def __init__(self, wcs, resolution=None):
        """
        Create a new WCS transform.  Resolution is the number of steps
        to interpolate between each input line segment to approximate its
        path.
        """
        Transform.__init__(self)
        self._resolution = resolution

        if wcs.naxis > 2:
            wcs = wcs.wcssub([1, 2]) 

        self.wcs = wcs
            
    def transform(self, ll):
        return self.wcs.world2pixel(ll)


    transform.__doc__ = Transform.transform.__doc__

    transform_non_affine = transform
    transform_non_affine.__doc__ = Transform.transform_non_affine.__doc__


    def inverted(self):
        return WcsPixel2SkyTransform(self.wcs, self._resolution)
    inverted.__doc__ = Transform.inverted.__doc__





class WcsPixel2SkyTransform(Transform):
    """
    The base Aitoff transform.
    """
    input_dims = 2
    output_dims = 2
    is_separable = False

    def __init__(self, wcs, resolution=None):
        """
        Create a new WCS transform.  Resolution is the number of steps
        to interpolate between each input line segment to approximate its
        path.
        """
        Transform.__init__(self)
        self._resolution = resolution

        if wcs.naxis > 2:
            wcs = wcs.wcssub([1, 2]) 

        self.wcs = wcs
            
    def transform(self, ll):
        return self.wcs.pixel2world(ll)


    transform.__doc__ = Transform.transform.__doc__

    transform_non_affine = transform
    transform_non_affine.__doc__ = Transform.transform_non_affine.__doc__


    def inverted(self):
        return WcsSky2PixelTransform(self.wcs, self._resolution)
    inverted.__doc__ = Transform.inverted.__doc__



class WcsSky2SkyTransform(Transform):
    """
    The base Aitoff transform.
    """
    input_dims = 2
    output_dims = 2
    is_separable = False

    def __init__(self, src_coord, dest_coord, resolution=None):
        """
        Create a new WCS transform.  Resolution is the number of steps
        to interpolate between each input line segment to approximate its
        path.
        """
        Transform.__init__(self)
        self._resolution = resolution

        self.src_coord = src_coord
        self.dest_coord = dest_coord

        self.coord_conv_func = coord_conv(src_coord, dest_coord)
        
        
    def transform(self, ll):
        lat1, lon1 = ll[:,0], ll[:,1]
        lat2, lon2 = self.coord_conv_func(lat1, lon1)
        
        return np.concatenate((lat2[:,np.newaxis], lon2[:,np.newaxis]),
                              1)


    transform.__doc__ = Transform.transform.__doc__

    transform_non_affine = transform
    transform_non_affine.__doc__ = Transform.transform_non_affine.__doc__


    def inverted(self):
        return WcsSky2SkyTransform(self.dest_coord, self.src_coord)
    inverted.__doc__ = Transform.inverted.__doc__







if 0:
    import pyfits
    import pywcs
    
    fname = "i013b4h0.fit"
    f = pyfits.open(fname)

    d, h = f[0].data, f[0].header
    d.shape = d.shape[-2:]
    
    _w = pywcs.WCS(h)
    w = _w.wcssub([1,2])
    
    x = array([0., 0.2, 0.4, 0.6])
    y = array([0., 0.2, 0.4, 0.6])
    xy = array([x,y]).transpose()
    
    radec = w.pixel2world(xy)

    from coords import GAL, FK5, FK4
    ss = WcsSky2SkyTransform(GAL, FK4)
    print ss.transform(array([[280.52288,-35.501269]]))
